package com.wuage.infrastructure.config;

import org.slf4j.MDC;
import org.springframework.cglib.proxy.Enhancer;
import org.springframework.cglib.proxy.MethodInterceptor;
import org.springframework.cglib.proxy.MethodProxy;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;

import java.lang.reflect.Method;
import java.util.Map;
import java.util.concurrent.Callable;
import java.util.concurrent.RejectedExecutionException;

/**
 * @author Shynin on 2019/3/28.
 */
@Configuration
public class ThreadPoolTaskExecutorConfig {

    @Bean("globalExecutor")
    public ThreadPoolTaskExecutor threadPoolTaskExecutor() {
        ThreadPoolTaskExecutor threadPoolTaskExecutor = create();
        threadPoolTaskExecutor.setCorePoolSize(20);
        threadPoolTaskExecutor.setMaxPoolSize(30);
        threadPoolTaskExecutor.setRejectedExecutionHandler((r, executor) -> {
                    try {
                        Thread e = new Thread(r, "Temporary task executor");
                        e.start();
                    } catch (Throwable var4) {
                        throw new RejectedExecutionException("Failed to start a new thread", var4);
                    }
                }
        );
        threadPoolTaskExecutor.setThreadNamePrefix("globalThreadPool-");
        return threadPoolTaskExecutor;
    }

    public ThreadPoolTaskExecutor create() {
        final Enhancer enhancer = new Enhancer();
        enhancer.setSuperclass(ThreadPoolTaskExecutor.class);
        enhancer.setCallback(new CglibProxy());
        return (ThreadPoolTaskExecutor) enhancer.create();
    }

    private class CglibProxy implements MethodInterceptor {


        private Callable handleCallable(Map<String, String> copyOfContextMap, Callable preCall) {
            return () -> {
                final Map<String, String> prev = MDC.getCopyOfContextMap();
                setContext(copyOfContextMap);
                try {
                    return preCall.call();
                } finally {
                    setContext(prev);
                }
            };
        }

        private Runnable handleRunnable(Map<String, String> copyOfContextMap, Runnable preRun) {
            return () -> {
                final Map<String, String> prev = MDC.getCopyOfContextMap();
                // setParent
                setContext(copyOfContextMap);
                try {
                    preRun.run();
                } finally {
                    //reset
                    setContext(prev);
                }
            };
        }

        private void setContext(Map<String, String> copyOfContextMap) {
            if (copyOfContextMap == null) {
                MDC.clear();
            } else {
                MDC.setContextMap(copyOfContextMap);
            }
        }


        @Override
        public Object intercept(Object o, Method method, Object[] objects, MethodProxy methodProxy) throws Throwable {
            if (method.getName().startsWith("submit") || method.getName().equals("execute")) {
                final Map<String, String> copyOfContextMap = MDC.getCopyOfContextMap();
                if (objects[0] instanceof Runnable) {
                    objects[0] = handleRunnable(copyOfContextMap, (Runnable) objects[0]);
                } else if (objects[0] instanceof Callable) {
                    objects[0] = handleCallable(copyOfContextMap, (Callable) objects[0]);
                }
            }
            return methodProxy.invokeSuper(o, objects);
        }


    }

}

